import os,time,sys,pdb
from pylab import plot,title,show,figure,savefig
from random import random
from math import sqrt,cos,sin,pi,acos,atan
from mpi4py import MPI
from numpy import array,dot,cross,isnan,linspace,empty,savez,c_,ravel, mgrid,save
from v3dsolver.vector import Vec
from v3dsolver.fluid_elements.fluid_elems import Saffman,FreeStream,FluidElem
from v3dsolver.boundary_manager import BoundaryManager,Geometry,Body
from v3dsolver.tree_manager.distribution_manager import DistributionManager,get_domain
from v3dsolver.tree_manager.vel_calculator import VelCalculator
from v3dsolver.tree_manager.tree_manager import TreeManager
from v3dsolver.advection_manager import AdvectionManager
from v3dsolver.stretching_manager import StretchingManager
from v3dsolver.diffusion_manager import DiffusionManager
from v3dsolver.merging_manager import MergingManager
from v3dsolver.annihilation_manager import AnnihilationManager
from v3dsolver.vorticity_interpolator import VorticityInterpolator
from enthought.tvtk.api import tvtk,write_data
from v3dsolver.writer import write

class Solver:
    """
    The overall 3D vortex solver, which combines functions like releasing blobs 
    from the boundary, advection, stretching, diffusion, merging, eating and 
    reflecting from boundary to simulate 3D flow past arbitrarily shaped
    solid bodies.
    """
    def __init__(self,bdry_file,rawdata_dir,blob_rad,delta,npmax,npseudo,dt,
                 freestream_val,dx_stretching,nu,R_bin,merge_frac,
                 cutoff_fraction):
        """
        Initialization of Solver. It initializes various parameters used in the
        simulation.
        
        bdry_file: filename of the vtk file representing the boundary
        rawdata_dir: directory path (relative to project directory) 
                     where simulation output will be stored in raw format
        blob_rad: radius of vortex blobs
        delta: distance of released blob from the boundary surface
        npmax: maximum no. of particles in a cell for it to be childless
        npseudo: no. of pseudo particles to be generated per cell
        dt: timestep of the simulation
        freestream_val: value of freestream velocity
        dx_stretching: distance between points on the same stretch stensil
        nu: kinematic viscosity
        R_bin: radius of binning (keep around 5*sqrt(nu*dt))
        merge_frac: merging radius as a fraction of cellsize = R_bin*sqrt(nu*dt)
        cutoff_fraction: fraction of avg vorticity below which a particle will
                         be discarded
        """
        self.geometry = Geometry([Body(bdry_file,blob_rad,delta)]) 
        self.rawdata_dir = os.path.join(os.path.abspath('../..'),
                                        rawdata_dir)
        self.delta = delta
        self.blob_rad = blob_rad
        self.npmax = npmax
        self.npseudo = npseudo
        self.dt = dt
        self.fs_val = freestream_val
        self.dx = dx_stretching
        self.nu = nu
        self.R_bin = R_bin
        self.merge_frac = merge_frac
        self.cutoff_frac = cutoff_fraction
        self.gridsize = 25
        self.gridpts, self.insidegridpts,self.totalgridpts = self.get_grid_pts2(self.gridsize)
    
    def get_grid_pts(self):
        gridpts = []
        r_list = linspace(1.,1.2,5)
        theta_list = linspace(0.,pi,15)
        phi_list = linspace(0.,2*pi,15)
        for r in r_list:
            for theta in theta_list:
                for phi in phi_list:
                    x = r*sin(theta)*cos(phi)
                    y = r*sin(theta)*sin(phi)
                    z = r*cos(theta)
                    pos = Vec(x,y,z)
                    newgridpt = Saffman(pos=pos,vel=Vec(),vor=Vec(),
                                        rad=self.blob_rad)
                    newgridpt.cause_or_effect = 1
                    gridpts += [newgridpt]
        return gridpts
        
    def get_grid_pts2(self,gridsize):
        gridpts,insidegridpts,totalgridpts = [],[],[]
        x_list = linspace(-2,2,gridsize)
        y_list = linspace(-2,2,gridsize)
        z_list = linspace(-2,2,gridsize)
        x_list, y_list, z_list = mgrid[-2:2:gridsize*1j, -2:2:gridsize*1j, -2:2:gridsize*1j]
        for i in range(len(x_list.flat)):
            pos = Vec(x_list.flat[i],y_list.flat[i],z_list.flat[i])
            newgridpt = Saffman(pos=pos,vel=Vec(),vor=Vec(),
                                rad=self.blob_rad)
            newgridpt.cause_or_effect = 1
            r = pos.length()
            if r < 1.:
                insidegridpts += [newgridpt]
            else:
                gridpts += [newgridpt]
            totalgridpts += [newgridpt]
        return gridpts,insidegridpts,totalgridpts
        
    def singlevortex_velfield(self):
        """
        Returns vel field due to a single vortex particle.
        """
        #elem = Saffman(pos=)
        pass
    
    def solve(self,tf,components='basdmnw'):
        """
        Solves flow till time tf.

        components: which components to solve for
                    (default: 'basdmn' which solves (b)oundary, (a)dvection, (s)tretching,
                     (d)iffusion, (m)erging and a(n)nihilation)
        """
        fs = FreeStream(self.fs_val)
        comm = MPI.COMM_WORLD

        if 'b' in components:
            bmanager = BoundaryManager(self.geometry,self.blob_rad,fs)
            this_outsideblobs = []
        else:
            self.totalelems = self.get_elems_random(n=2)
            dm = DistributionManager(comm,self.totalelems)
            this_outsideblobs = dm.get_this_proc_elems()        
        this_bblobs = []
        this_rblobs = []
        
        t = 0
        count = 0
        while t<tf:        
            print '\nt = '+str(t)+'\n'+'='*len('t = '+str(t))
            
            # Release blobs from the boundary
            if 'b' in components:
                time1 = time.time()
                this_bblobs,this_rblobs,this_new_outsideblobs = bmanager.get_released_particles(this_outsideblobs)
                time2 = time.time()
                this_outsideblobs += this_new_outsideblobs
                
                print 'time taken by boundary:' + str(time2 - time1)
                print 'this_outsideblobs after boundary: '+str(len(this_outsideblobs))
                total_vor = Vec()
                for blob in this_outsideblobs+this_bblobs:
                    total_vor += blob.vor
                print 'Total vor:', total_vor

            # Advect all elements
            if 'a' in components:
                time1 = time.time()
                advection_manager = AdvectionManager(this_outsideblobs+this_bblobs+this_rblobs,
                                    self.dt,fs,self.npmax,self.npseudo,self.blob_rad,'p')
                advection_manager.move_RK2()
                time2 = time.time()
                print 'time taken by convection:' + str(time2-time1)
                print 'this_outsideblobs after advection: '+str(len(this_outsideblobs))
                avg_vel = Vec()
                avg_pos = Vec()
                for elem in this_outsideblobs:
                    avg_vel += elem.vel
                    avg_pos += elem.pos
                avg_vel = avg_vel/len(this_outsideblobs)
                avg_pos = avg_pos/len(this_outsideblobs)
                print 'Avg_vel:', avg_vel
                print 'Avg_pos:', avg_pos
            
            # Save arrays for post processing
            if 'w' in components:
                # computing velocities
                elems = this_outsideblobs + self.gridpts + this_bblobs + this_rblobs
                domain = get_domain(elems,findglobal=True) # global domain calculation
                tree = TreeManager(elems,domain,npmax=5,npseudo=16,rad=self.blob_rad,
                                   comm=comm)
                tree.root.effect_elems = self.gridpts
                #VelCalculator(tree).parallel_vel_computation()
                for elem in tree.root.effect_elems:
                    elem.vel += fs.get_vel(elem)
                
                # writing blob data
                blobfile = os.path.join(self.rawdata_dir,
                                str(comm.rank)+'b_dt='+str(self.dt)+'_'+str(count))
                write(this_outsideblobs,blobfile)
                print 'written blob data'
                
                # writing grid data
                #gridfile = os.path.join(self.rawdata_dir,
                #                str(comm.rank)+'g_t='+str(round(t,4)))
                #write(self.totalgridpts,gridfile,isgrid=True,gridsize=self.gridsize)
                #print 'written grid data'
            
            # Stretch vorticities of all elements
            if 's' in components:
                initvor = Vec()
                for blob in this_outsideblobs:
                    initvor += blob.vor
                stretching_manager = StretchingManager(this_outsideblobs+this_bblobs+this_rblobs,
                                     self.dx,self.dt,fs,self.npmax,
                                     self.npseudo,self.blob_rad,'p')
                stretching_manager.update_vors()
                finalvor = Vec()
                for blob in this_outsideblobs:
                    finalvor += blob.vor
                print 'this_outsideblobs after stretching: '+str(len(this_outsideblobs))
                print 'total vorticity before stretching:', initvor
                print 'total vorticity after stretching:', finalvor
                pdb.set_trace()
            
            # Diffuse all elements
            if 'd' in components:
                initvor = Vec()
                initvormom = Vec()
                for blob in this_outsideblobs:
                    initvor += blob.vor
                    initvormom += blob.vor.cross(blob.pos)
                time1 = time.time()
                diffusion_manager = DiffusionManager(this_outsideblobs,self.nu,self.dt,
                                    self.R_bin,self.merge_frac,mode='p',geometry=self.geometry)
                diffusion_manager.preprocessing()
                diffusion_manager.diffuse()
                time2 = time.time()
                finalvor = Vec()
                finalvormom = Vec()
                for blob in this_outsideblobs:
                    finalvor += blob.vor
                    finalvormom += blob.vor.cross(blob.pos)
                print 'total vorticity before diffusion:', initvor
                print 'total vorticity after diffusion:', finalvor
                print 'total vorticity momentum before diffusion:', initvormom
                print 'total vorticity momentum after diffusion:', finalvormom
                print 'time taken by diffusion:'+str(time2-time1)
                print 'this_outsideblobs after diffusion: '+str(len(this_outsideblobs))
                
            # Annihilate particles which are going to go inside the boundary
            if 'n' in components:
                maxvor = max([elem.vor.length() for elem in this_outsideblobs]) # calculating max of rblobs
                print 'minvor:', maxvor
                annihilation_manager = AnnihilationManager(this_outsideblobs,self.cutoff_frac)
                annihilation_manager.annihilate(maxvor)
                print 'this_outsideblobs after annihilation: '+str(len(this_outsideblobs))
                
            # Merge particles and reflect them from the boundary
            if 'm' in components:
                time1 = time.time()
                diffusion_manager.reflect_from_bdry()
                time2 = time.time()
                print 'time taken by reflection:'+str(time2-time1)
                
                time1 = time.time()
                diffusion_manager.merge()
                time2 = time.time()
                finalvor2 = Vec()
                finalvormom2 = Vec()
                for blob in this_outsideblobs:
                    finalvor2 += blob.vor
                    finalvormom2 += blob.vor.cross(blob.pos)
                print 'total vorticity after merging and reflection:', finalvor2
                print 'total vorticity after merging and reflection:', finalvormom2
                print 'time taken by merging:'+str(time2-time1)
                print 'this_outsideblobs after merging and reflection: '+ str(len(this_outsideblobs))
            
            # Update time
            t += self.dt
            count += 1
            #pdb.set_trace()
                
    def get_elems(self):
        """
        Returns elements
        """
        n = 500
        elems = []
        for i in range(n):
            x = cos((2.*pi*i)/n)
            y = sin((2.*pi*i)/n)
            z = 0.
            vor = cross(array([x,y,z]),array([0.,0.,-0.01]))
            elems += [Saffman(pos=array([x,y,z]),vor=vor,
                              vel=array([0.,0.,0.]),rad=self.blob_rad)]
        n = 0#100
        for i in range(n):
            x = cos((2.*pi*i)/n)
            y = sin((2.*pi*i)/n)
            z = 0.2
            vor = cross(array([x,y,z]),array([0.,0.,-0.01]))
            elems += [Saffman(pos=array([x,y,z]),vor=vor,
                              vel=array([0.,0.,0.]),rad=self.blob_rad)]
        return elems
        
    def get_elems_random(self,n):
        """
        Returns a list of n randomly selected particles within the 1x1x1 cube.
        """
        elems = []
        rad = 1.
        for i in range(n):
            x = rad*random()
            y = rad*random()
            z = rad*random()
            elems += [Saffman(pos=Vec(x,y,z),vor=Vec(1.,1.,1.),
                                   vel=Vec(),rad=self.blob_rad)]
        return elems

if __name__=="__main__":
    casename = sys.argv[1]
    Re = 20.
    nu = 1./Re
    vortexsolver = Solver(bdry_file=casename+'.vtk',rawdata_dir='examples/'+casename+'/raw/Re10_vrt_corrected',
                          blob_rad=0.1,delta=0.1,npmax=5,npseudo=48,dt=0.01,
                          freestream_val=Vec(0.,0.,1.),
                          dx_stretching=0.001,nu=nu,R_bin=sqrt(17),
                          merge_frac=0.4,cutoff_fraction=0.005)
    try:
        vortexsolver.solve(tf=1,components='d')
    except Exception as e:
        import sys
        exc_type,exc_value,exc_tb = sys.exc_info()
        import traceback
        print traceback.print_exc()
        pdb.set_trace()
